#include <cuda_runtime.h>
#include <stdio.h>
#include <math.h>

// CUDA 错误检查宏
#define CUDA_CHECK(call)                                                        \
    do                                                                          \
    {                                                                           \
        cudaError_t err = call;                                                 \
        if (err != cudaSuccess)                                                 \
        {                                                                       \
            fprintf(stderr, "CUDA error in %s (%s:%d): %s\n",                   \
                    __FUNCTION__, __FILE__, __LINE__, cudaGetErrorString(err)); \
            exit(EXIT_FAILURE);                                                 \
        }                                                                       \
    } while (0)

// 每个线程处理的输出元素数量
#define ELEMENTS_PER_THREAD 4

// 双线性上采样内核优化版
__global__ void bilinear_upsample_kernel_optimized(
    int IC, int IH, int IW, int OH, int OW,
    const float *__restrict__ inp,
    float *__restrict__ out)
{
    // 计算每个线程块中线程的全局索引
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total = IC * OH * OW;

    // 计算每个线程需要处理的输出元素的起始索引
    int start_idx = tid * ELEMENTS_PER_THREAD;

    // 使用循环让每个线程处理多个输出元素
    for (int i = 0; i < ELEMENTS_PER_THREAD; ++i)
    {
        int out_idx = start_idx + i;
        if (out_idx >= total)
            continue;

        // 解码全局索引为通道、高度和宽度索引
        // 此线程处理的通道索引
        int c = out_idx / (OH * OW);
        int remainder = out_idx % (OH * OW);
        int oh = remainder / OW;
        int ow = remainder % OW;

        // 计算缩放因子
        float scale_h = (IH - 1) / (float)(OH - 1);
        float scale_w = (IW - 1) / (float)(OW - 1);

        // 映射到输入空间
        float ih_f = oh * scale_h;
        float iw_f = ow * scale_w;

        // 确定输入空间的四个邻近像素
        int ih0 = floorf(ih_f);
        int iw0 = floorf(iw_f);
        int ih1 = ih0 + 1 < IH ? ih0 + 1 : ih0;
        int iw1 = iw0 + 1 < IW ? iw0 + 1 : iw0;

        // 计算小数部分
        float delta_h = ih_f - ih0;
        float delta_w = iw_f - iw0;

        // 计算双线性插值的权重
        float w1 = (1.0f - delta_h) * (1.0f - delta_w);
        float w2 = (1.0f - delta_h) * delta_w;
        float w3 = delta_h * (1.0f - delta_w);
        float w4 = delta_h * delta_w;

        // 计算输入数据的基地址
        int inp_base = c * IH * IW;

        // 获取四个邻近像素的值
        float val1 = inp[inp_base + ih0 * IW + iw0];
        float val2 = inp[inp_base + ih0 * IW + iw1];
        float val3 = inp[inp_base + ih1 * IW + iw0];
        float val4 = inp[inp_base + ih1 * IW + iw1];

        // 计算插值后的输出值
        out[out_idx] = w1 * val1 + w2 * val2 + w3 * val3 + w4 * val4;
    }
}

// 主机函数调用优化后的双线性上采样内核
void bilinear_upsample_optimized(
    int IC, int IH, int IW, int OH, int OW,
    const float *inp_device,
    float *out_device)
{
    // 计算总输出元素数
    int total = IC * OH * OW;

    // 计算每个线程处理的输出元素数量
    int elements_per_thread = ELEMENTS_PER_THREAD;

    // 计算需要的线程数量
    int threads = (total + elements_per_thread - 1) / elements_per_thread;

    // 定义每个线程块的线程数
    int threads_per_block = 256;
    int num_blocks = (threads + threads_per_block - 1) / threads_per_block;

    // 启动优化后的CUDA内核
    bilinear_upsample_kernel_optimized<<<num_blocks, threads_per_block>>>(
        IC, IH, IW, OH, OW, inp_device, out_device);

    // 检查内核执行是否成功
    CUDA_CHECK(cudaGetLastError());
    CUDA_CHECK(cudaDeviceSynchronize());
}
